import numpy as np
from scipy import special
from chaco.api import *
from traits.api import *
from traitsui.api import *
from chaco.tools.api import PanTool, ZoomTool
from enable.api import ComponentEditor
import time
from PyQt5.QtCore import QThread

from mayavi.core.api import Engine, PipelineBase
from mayavi.core.ui.api import MayaviScene, MlabSceneModel, SceneEditor
from tvtk.api import tvtk

engine = Engine()
engine.start()


class Form(HasTraits):
    frequencies = List(Float, [3, 10])
    elevations = List(Float, [0, 10])
    azimuthStart = Float(0)
    azimuthStop = Float(45)
    azimuthNumber = Int(91)
    plot = Instance(Plot)
    thread = QThread()

    def calculate(self):
        v = self.v
        f = self.f
        fr = .3
        k = 2 * np.pi / .3 * fr
        v3 = np.concatenate((v[f[:, 0]][:, np.newaxis], v[f[:, 1]][:, np.newaxis], v[f[:, 2]][:, np.newaxis]), 1)
        l = np.linalg.norm(np.roll(v3, -1, 1) - v3, axis=-1)
        c = v3.mean(1)
        r = np.linalg.norm(c[:, np.newaxis] - v3, axis=-1)
        t = (r + np.roll(r, -1, -1) + l) / 2
        t = np.sqrt(t * (t - r) * (t - np.roll(r, -1, -1)) * (t - l))
        a = t.sum(-1)
        t = 2 * t / l
        fs = (t * np.log((r + np.sqrt(r ** 2 - t ** 2)) * (np.roll(r, -1, -1) + np.sqrt(np.roll(r, -1, -1) ** 2 - t ** 2)) / t ** 2)).sum(-1) / a - 1j * k
        ef = []
        et = []
        for i in range(f.shape[0] - 1):
            fi = f[i]
            for j in range(i + 1, f.shape[0]):
                fj = f[j]
                if len(set(fj) - set(fi)) == 1:
                    ef.append([i, j])
                    et.append([list(set(fi) - set(fj))[0], list(set(fj) - set(fi))[0]])
        ef = np.array(ef)
        et = np.array(et)
        ei = np.array([1, 0, 0])
        inc = np.array([0, 0, -1])
        theta = np.arange(181)
        costheta = special.cosdg(theta)
        sintheta = special.sindg(theta)
        z = 0
        rh = 0
        for i in [0, 1]:
            for j in [0, 1]:
                t = np.repeat(fs[ef[:, j]][:, np.newaxis], ef.shape[0], 1).astype(complex)
                n = ef[:, i][:, np.newaxis] != ef[:, j][np.newaxis]
                r = np.linalg.norm(c[ef[:, i]][:, np.newaxis] - c[ef[:, j]][np.newaxis], axis=-1)[n]
                t[n] = np.exp(-1j * k * r) / r
                z += (np.inner(c[ef[:, i]] - v[et[:, i]], c[ef[:, j]] - v[et[:, j]]) / 4 - 1 / k ** 2) * t * (-1) ** (i + j)
            rh += -(c[ef[:, i]] - v[et[:, i]]) @ ei * np.exp(-1j * k * c[ef[:, i]] @ inc) / 2 * (-1) ** i
        u = np.linalg.solve(z, rh)
        current = np.zeros(f.shape[0], dtype=complex)
        np.add.at(current,ef[:, 0],u)
        np.add.at(current,ef[:, 1],u)
        self.current = np.abs(current)
        sEplane = np.column_stack((sintheta, np.zeros(theta.size), costheta))
        sHplane = np.column_stack((np.zeros(theta.size), -sintheta, costheta))
        esEplane = np.column_stack((costheta, np.zeros(theta.size), -sintheta))
        esHplane = np.column_stack((np.ones(theta.size), np.zeros(theta.size), np.zeros(theta.size)))

        rcsEplane = np.sqrt(np.pi) * np.inner(esEplane, c[ef[:, 0]] - v[et[:, 0]]) * np.exp(1j * k * np.inner(sEplane, c[ef[:, 0]])) @ u \
                    - np.sqrt(np.pi) * np.inner(esEplane, c[ef[:, 1]] - v[et[:, 1]]) * np.exp(1j * k * np.inner(sEplane, c[ef[:, 1]])) @ u
        rcsHplane = np.sqrt(np.pi) * np.inner(esHplane, c[ef[:, 0]] - v[et[:, 0]]) * np.exp(1j * k * np.inner(sHplane, c[ef[:, 0]])) @ u \
                    - np.sqrt(np.pi) * np.inner(esHplane, c[ef[:, 1]] - v[et[:, 1]]) * np.exp(1j * k * np.inner(sHplane, c[ef[:, 1]])) @ u

        n = np.arange(1, 81)
        ka = k
        pds = special.lpmv(1, n, costheta[1:-1, np.newaxis]) / sintheta[1:-1, np.newaxis]
        pdd = special.lpmv(2, n, costheta[1:-1, np.newaxis]) + pds * costheta[1:-1, np.newaxis]
        pds = np.vstack((-n * (n + 1) / 2, pds, (-1) ** n * n * (n + 1) / 2))
        pdd = np.vstack((-n * (n + 1) / 2, pdd, -(-1) ** n * n * (n + 1) / 2))
        jn = special.jv(n + .5, ka)
        hn = special.hankel2(n + .5, ka)
        jp = special.jv(n - .5, ka) - n / ka * special.jv(n + .5, ka)
        hp = special.hankel2(n - .5, ka) - n / ka * special.hankel2(n + .5, ka)
        ae = 2j * (-1.) ** n * (2 * n + 1) / n / (n + 1) / ka * jn / hn * np.sqrt(np.pi)
        am = 2j * (-1.) ** n * (2 * n + 1) / n / (n + 1) / ka * jp / hp * np.sqrt(np.pi)
        rcsEplane0 = pds @ ae - pdd @ am
        rcsHplane0 = pdd @ ae - pds @ am

        self.plot.data.update_data(
            x=theta,
            rcsEplane=20 * np.log10(np.abs(rcsEplane)),
            rcsHplane=20 * np.log10(np.abs(rcsHplane)),
            rcsEplane0=20 * np.log10(np.abs(rcsEplane0)),
            rcsHplane0=20 * np.log10(np.abs(rcsHplane0)),
        )
        # mlab = self.polydata.mlab
        # points = np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0], [1, 1, 0], [2, 0, 0]], 'f')
        # triangles = np.array([[0, 1, 2], [1, 3, 2], [1, 4, 3], ])
        # scalars = np.array([0, .1, .2, .3, .4], dtype=float)
        # mlab.triangular_mesh(points[:, 0], points[:, 1], points[:, 2], triangles, scalars=scalars)

    def traits_init(self):
        np.seterr(divide='ignore')
        with open('sphere0.25m.dat') as file:
            s = np.array(file.read().split())
        self.v = s[np.argwhere(s == 'GRID*') + [2, 3, 7]].astype(float)
        self.f = s[np.argwhere(s == 'CTRIA3') + [3, 4, 5]].astype(int) - 1
        self.thread.run = self.calculate
        self.plot = Plot(ArrayPlotData())
        self.plot.data.update_data(
            x=[],
            rcsEplane=[],
            rcsHplane=[],
            rcsEplane0=[],
            rcsHplane0=[],
        )
        self.plot.x_axis.title = "Azimuth/deg"
        self.plot.y_axis.title = "RCS/dBsm"
        self.plot.plot(("x", "rcsEplane"), name='RCS E plane', color='auto')
        self.plot.plot(("x", "rcsHplane"), name='RCS H plane', color='auto')
        self.plot.plot(("x", "rcsEplane0"), name='RCS E plane Mie', color='auto')
        self.plot.plot(("x", "rcsHplane0"), name='RCS H plane Mie', color='auto')
        legend = Legend()
        legend.plots = self.plot.plots
        self.plot.overlays.append(legend)
        self.plot.tools.append(
            PanTool(self.plot, constrain_key="shift")
        )
        self.plot.overlays.append(
            ZoomTool(component=self.plot, tool_mode="box", always_on=False)
        )

    button = Button('计算')
    button1 = Button('显示表面电流')

    def _button_fired(self):
        self.calculate()
    def _button1_fired(self):
        self.update_polydata()

    polydata = Instance(MlabSceneModel, args={'engine': engine})

    # @observe('polydata.activated')
    # def update_polydata(self, event=None):
    #     mlab = self.polydata.mlab
    #     v = self.v
    #     f = self.f
    #     scalars = np.arange(v.shape[0])
    #     mlab.triangular_mesh(v[:, 0], v[:, 1], v[:, 2], f, scalars=self.current)
    #

    @observe('polydata.activated')
    def update_polydata(self, event=None):
        v = self.v
        f = self.f
        mlab = self.polydata.mlab
        mesh = tvtk.PolyData(points=v, polys=f)
        mesh.cell_data.scalars = self.current
        surf = mlab.pipeline.surface(mesh)
        mlab.pipeline.surface(mlab.pipeline.extract_edges(surf), color=(0, 0, 0), )

    traits_view = View(
        Tabbed(
            Group(
                UItem(
                    "plot",
                    editor=ComponentEditor(),
                    resizable=True,
                ),
                HGroup(
                    UItem('button'),
                    spring
                )
            ),
            Group(
                UItem(
                    "polydata",
                    editor=SceneEditor(scene_class=MayaviScene),
                    resizable=True,
                ),
                HGroup(
                    UItem('button1'),
                    spring
                ),
            ),
        ),
        width=1200,
        height=900,
        resizable=True,

    )


    parameter_view = View(
        Group(
            UItem('frequencies'),
            label='频率/GHz',
            show_border=True,
        ),
        Group(
            UItem('elevations'),
            label='高低角/°',
            show_border=True,
        ),
        Item('azimuthStart', label='方位角起始/°'),
        Item('azimuthStop', label='方位角终止/°'),
        Item('azimuthNumber', label='方位角数量'),
        height=600,
        resizable=True,
        buttons=OKCancelButtons,
    )

if __name__ == '__main__':
    Form().configure_traits()
